package com.chujian.framework.mybatis.plugins;

import cn.hutool.core.collection.CollUtil;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.chujian.framework.mybatis.auth.QueryDataAuth;
import com.chujian.framework.mybatis.auth.QueryDataAuthHolder;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectExpressionItem;
import net.sf.jsqlparser.statement.select.SelectItem;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.io.StringReader;
import java.sql.SQLException;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author ：chujian
 * @since : 2021-8-12
 * 数据权限拦截器
 */
@Slf4j
public class QueryDataAuthInterceptor implements InnerInterceptor {

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {

        QueryDataAuth queryDataAuth = QueryDataAuthHolder.getQueryDataAuth();
        if (null == queryDataAuth) return;
        if (StringUtils.isEmpty(queryDataAuth.getWhereSql()) && CollUtil.isEmpty(queryDataAuth.getDenyCols())) return;

        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        String newSql = getAuthedSql(mpBs.sql(), queryDataAuth);
        mpBs.sql(newSql);
    }

    public String getAuthedSql(String originalSql, QueryDataAuth queryDataAuth) throws SQLException {
        try {
            // JSqlParser解析原sql
            CCJSqlParserManager parserManager = new CCJSqlParserManager();
            Select select = (Select) parserManager.parse(new StringReader(originalSql));
            PlainSelect selectBody = (PlainSelect) select.getSelectBody();

            // where权限注入
            if (StringUtils.isNotEmpty(queryDataAuth.getWhereSql())) {
                Expression authExpress = CCJSqlParserUtil.parseCondExpression(queryDataAuth.getWhereSql());
                if (selectBody.getWhere() == null || StringUtils.isEmpty(selectBody.getWhere().toString())) {
                    selectBody.setWhere(new Parenthesis(authExpress));
                } else {
                    selectBody.setWhere(new AndExpression(selectBody.getWhere(), new Parenthesis(authExpress)));
                }
            }

            // 列权限注入
            if (CollUtil.isNotEmpty(queryDataAuth.getDenyCols())) {
                List<SelectItem> selectItems = selectBody.getSelectItems()
                        .stream()
                        .filter(selectItem -> {
                            if (selectItem instanceof SelectExpressionItem) {
                                SelectExpressionItem item = (SelectExpressionItem) selectItem;
                                if (item.getExpression() instanceof Column) {
                                    Column column = (Column) item.getExpression();
                                    String columnName = column.getColumnName();
                                    return !queryDataAuth.getDenyCols().contains(columnName.toLowerCase());
                                }
                            }
                            return true;
                        }).collect(Collectors.toList());
                selectBody.setSelectItems(selectItems);
            }
            return selectBody.toString();
        } catch (JSQLParserException e) {
            log.error("设置权限失败", e);
        }
        return originalSql;
    }

}
